from AStock.ASUtil import tushare_client
from AStock.ASSetting import settings
from pymongo import UpdateOne
import argparse


def _main_(args):
    begin_date = args.begin_date
    end_date = args.end_date

    pro = tushare_client()

    # 交易日历
    coll = settings.database.trade_cal
    df_trade_cal = pro.trade_cal(exchange='', start_date=begin_date, end_date=end_date)
    bulk = []
    for tup in df_trade_cal.itertuples():
        update_one = UpdateOne(
            {'cal_date': int(tup.cal_date)},
            {'$set': {'cal_date': int(tup.cal_date), 'is_open': tup.is_open}},
            upsert=True
        )
        bulk.append(update_one)
    try:
        coll.bulk_write(bulk)
        print('saved {} items to trade_cal'.format(len(bulk)))
    except Exception as e:
        print(e)


if __name__ == '__main__':
    argparser = argparse.ArgumentParser(description='save trade calendar')
    argparser.add_argument('-b', '--begin-date', required=True,
                           help='trade calendar begin date, yyyymmdd format')
    argparser.add_argument('-e', '--end-date', required=True,
                           help='trade calendar end date, yyyymmdd format')
    args = argparser.parse_args()
    _main_(args)
